﻿using Microsoft.Extensions.Logging;
using System.Net.Sockets;
using System.Text;
using System.Text.RegularExpressions;

namespace Away.NetMap.Nmap;

public class ServiceScannerParams
{
    public string IP { get; set; } = null!;
    public int Port { get; set; }
    /// <summary>
    /// 协议： TCP|UDP
    /// </summary>
    public string Protocol { get; set; } = "TCP";
    /// <summary>
    /// 线程数
    /// </summary>
    public int Threads { get; set; } = 100;
    /// <summary>
    /// 超时时间
    /// </summary>
    public int Timeout { get; set; } = 300;
}

[ServiceInject(ServiceLifetime.Scoped, true)]
public class ServiceScannerService
{
    private readonly IServiceScope _scope;
    public ServiceScannerService(IServiceProvider serviceProvider)
    {
        _scope = serviceProvider.CreateScope();
        ScanCompleted += Completed;
    }

    private IServiceProvider ServiceProvider => _scope.ServiceProvider;
    private PortRepository Rep => ServiceProvider.GetRequiredService<PortRepository>();
    private ServiceProbeFile ServiceProbe => ServiceProvider.GetRequiredService<ServiceProbeFile>();
    private ILogger<ServiceScannerService> Logger => ServiceProvider.GetRequiredService<ILogger<ServiceScannerService>>();

    private readonly CancellationTokenSource _source = new();
    private ServiceScannerParams? _params;
    public void Run(ServiceScannerParams p, CancellationToken token)
    {
        token.Register(_source.Cancel);
        System.Diagnostics.Stopwatch stopwatch = new();
        stopwatch.Start();

        _params = p;
        Logger.LogInformation("开始探测服务：{}://{}:{} ", p.Protocol, p.IP, p.Port);
        var threads = p.Threads;

        var probes = ServiceProbe.Probes.Where(o => o.Protocol == _params.Protocol).OrderBy(o => o.Rarity);
        if (probes == null)
        {
            return;
        }
        try
        {
            probes
            .AsParallel()
            .OrderByDescending(o => o)
            .WithDegreeOfParallelism(threads)
            .WithCancellation(_source.Token)
            .ForAll(ScanProbeOne);
        }
        catch (OperationCanceledException)
        {

        }

        Logger.LogInformation("结束探测服务：{}://{}:{} 耗时：{}", p.Protocol, p.IP, p.Port, stopwatch.Elapsed);
    }

    private void ScanProbeOne(ServiceProbe probe)
    {
        try
        {
            ProtocolType protocol = _params!.Protocol == "TCP" ? ProtocolType.Tcp : ProtocolType.Udp;
            using Socket socket = new(AddressFamily.InterNetwork, SocketType.Stream, protocol);
            socket.Connect(_params!.IP, _params!.Port);

            if (probe.TotalWaitms > 0)
            {
                socket.ReceiveTimeout = probe.TotalWaitms;
            }
            if (probe.TcpWrappedms > 0)
            {
                socket.SendTimeout = probe.TcpWrappedms;
            }

            if (!string.IsNullOrWhiteSpace(probe.ProbeString))
            {
                socket.Send(Encoding.Default.GetBytes(probe.ProbeString));
            }

            byte[] buffer = new byte[socket.ReceiveBufferSize];
            int len = socket.Receive(buffer);
            var context = Encoding.UTF8.GetString(buffer[0..len]);

            if (string.IsNullOrWhiteSpace(context))
            {
                return;
            }
            Logger.LogTrace("发送探针{}：\r\n{} \r\n探针获取内容：\r\n{}", probe.ProbeName, probe.ProbeString, context);
            ScanMoreProtocol(context, probe.Matches);
        }
        catch (OperationCanceledException)
        {

        }
        catch (SocketException ex)
        {
            Logger.LogTrace(ex, "扫描探针失败：{}://{}:{}", _params!.Protocol, _params.IP, _params.Port);
        }
    }

    private void ScanMoreProtocol(string context, List<ProbeMatchItem> matches)
    {
        matches
            .AsParallel()
            .OrderByDescending(o => o)
            .WithDegreeOfParallelism(_params!.Threads)
            .WithCancellation(_source.Token)
            .ForAll(o => ScanOneProtocol(context, o));
    }

    private void ScanOneProtocol(string context, ProbeMatchItem match)
    {
        var flag = ScanOneProtocol(context, match as ProbeSoftMatchItem);
        if (flag)
        {
            return;
        }

        if (match.SoftMatches.Count > 0)
        {
            ScanMoreProtocol(context, match.SoftMatches);
            return;
        }
    }

    private void ScanMoreProtocol(string context, List<ProbeSoftMatchItem> matches)
    {
        try
        {
            matches
                .AsParallel()
                .OrderByDescending(o => o)
                .WithDegreeOfParallelism(_params!.Threads)
                .WithCancellation(_source.Token)
                .ForAll(o => ScanOneProtocol(context, o));
        }
        catch (OperationCanceledException)
        {
        }
    }

    private bool ScanOneProtocol(string context, ProbeSoftMatchItem match)
    {
        var reg = Regex.Match(context, match.Pattern);
        if (!reg.Success)
        {
            return false;
        }
        _source.Cancel();
        Completed(match.Service, context);
        return true;
    }


    private delegate void ScanCompletedEventHandler(string service, string context);
    private event ScanCompletedEventHandler ScanCompleted;
    private static readonly object _lock = new();
    private void Completed(string service, string context)
    {
        lock (_lock)
        {
            Logger.LogTrace("{}://{}:{}", service, _params!.IP, _params.Port);
            Rep.Update(_params.IP, _params.Port, service);
        }
    }
}